'''
	Author：Yassin
	用于选取股票价格数据，包含函数 Sel_from_daily_ret, get_trade_dates
'''

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def Get_trade_dates(RangeL : str, RangeR:str) -> pd.DataFrame:
    '''
        获取交易日历
        RangeL，RangeR :str 获取交易日历的时间范围, format = YYYYMMDD
        返回一个DataFrame，包括交易日 cal_date 和其上一个交易日 pretrade_date
    '''
    df = pd.read_csv("~/Desktop/mpacc2/data_invest/temp/data_lib/trade_date.csv")
    df = pd.DataFrame(df)
    df["cal_date"] = pd.to_datetime(df["cal_date"], format = "%Y%m%d")
    df["pretrade_date"] = pd.to_datetime(df["pretrade_date"], format = "%Y%m%d")
    df = df.sort_values("cal_date", ascending = True)
    df = df.reset_index(drop = True)
    RangeL = pd.to_datetime(RangeL, format = "%Y%m%d")
    RangeR = pd.to_datetime(RangeR, format = "%Y%m%d")
    df = df[df["cal_date"] >= RangeL]
    df = df[df["cal_date"] <= RangeR]
    df = df.reset_index(drop = True)
    return df


Read_flag = False
Read_df = pd.DataFrame()

def Sel_from_daily_ret(RangeL : str, RangeR : str, Stock_code : list = [], Markettype:list = [1,4]):
    '''
        获取交易行情数据
        RangeL，RangeR : str 获取交易日历的时间范围, format = YYYYMMDD
        Stock_code : list ，为所需要行情的股票代码（int），默认为空，即查询全部股票的行情数据
        Markettype : list ，为所需的市场，默认为 1， 为主板市场
        返回一个DataFrame， 为所需要时间范围内，所需要股票的行情数据, 包括 开高低收，交易量和交易额，流通市值和总市值，所在交易市场及交易状态
    '''
    global Read_df
    global Read_flag
    if Read_flag == False:
        data = pd.read_hdf(r"~/Desktop/mpacc2/data_invest/temp/data_lib/returnHdf.h5")
        Read_df = data
        Read_flag = True
    else:
        data = Read_df

    target_L = ['Stkcd', 'Trddt', 'Opnprc', 'Hiprc', 'Loprc', 'Clsprc', 'Dnshrtrd', 'Dnvaltrd', 'Dsmvosd', 'Dsmvtll', 'Markettype', 'Trdsta', "Dretwd"]
    data = data[target_L].reset_index(drop = True)

    if len(Stock_code) != 0:
        target_stock = []
        for item in Stock_code:
            target_stock.append(data[data["Stkcd"] == item])
        data = pd.concat(target_stock)

    target_market = []
    for item in Markettype:
        target_market.append(data[data["Markettype"] == item])
    df = pd.concat(target_market).reset_index(drop = True)

    df["Trddt"] = pd.to_datetime(df["Trddt"])
    df = df.sort_values("Trddt", ascending = True).reset_index(drop = True)
    RangeL = pd.to_datetime(RangeL)
    RangeR = pd.to_datetime(RangeR)
    df = df[df["Trddt"] >= RangeL]
    df = df[df["Trddt"] <= RangeR]
    df = df.reset_index(drop = True)
    return df

if __name__ == "__main__":
    print(Sel_from_daily_ret("20150101","20150201",Stock_code=[600519]))
    df = Sel_from_daily_ret("20150101","20150201",Stock_code=[600519])
    df.to_csv("test.csv")